# -*- coding: utf-8 -*-
from __future__ import print_function

import base64
import subprocess
import argparse
import agent_config
import os
import json
from PIL import Image

from openai import OpenAI

from utils import sr_utils, sr_s3

"""
5010: "内部脚本执行超时",
5011: "输入文件或水印文件不存在",
5012: "提取缩略图失败",
5013: "添加水印失败",
5014: "提取帧失败"
5015: "NSFW检测不通过"
"""


def extract_thumbnail(input_video, output_image='thumbnail.jpg', time_second=0):
    """从视频中提取指定时间点的帧作为缩略图，默认提取第一秒第一帧"""
    # 构建ffmpeg命令
    cmd = [
        'ffmpeg', '-i', input_video,
        '-ss', '00:00:{}'.format(time_second),
        '-vframes', '1',
        '-y',
        output_image
    ]

    # 执行命令
    try:
        subprocess.check_output(cmd, stderr=subprocess.STDOUT, universal_newlines=True)
    except subprocess.CalledProcessError as e:
        print("Failed to extract thumbnail: {}".format(e.output))
        return None
    print("Thumbnail saved to: {}".format(output_image))
    return output_image


def extract_frame(input_video, output_path, time_second, max_size=720):
    """提取视频指定时间点的帧并压缩至最大边不超过max_size"""
    # 临时文件路径
    temp_output = output_path + '.temp.png'

    # 提取帧
    cmd = [
        'ffmpeg', '-i', input_video,
        '-ss', '00:00:{}'.format(time_second),
        '-vframes', '1',
        '-y',
        temp_output
    ]

    try:
        subprocess.check_output(cmd, stderr=subprocess.STDOUT, universal_newlines=True)
    except subprocess.CalledProcessError as e:
        print("Failed to extract frame at {} seconds: {}".format(time_second, e.output))
        # 清理可能创建的临时文件
        if os.path.exists(temp_output):
            os.remove(temp_output)
        return False

    # 检查临时文件是否存在
    if not os.path.exists(temp_output):
        print("Temporary frame file was not created")
        return False

    # 压缩图片
    try:
        with Image.open(temp_output) as img:
            # 计算压缩比例
            width, height = img.size
            if max(width, height) > max_size:
                ratio = max_size / max(width, height)
                new_size = (int(width * ratio), int(height * ratio))
                img = img.resize(new_size, Image.LANCZOS)
            img.save(output_path)
    except Exception as e:
        print("Failed to process image: {}".format(str(e)))
        # 清理文件
        if os.path.exists(temp_output):
            os.remove(temp_output)
        return False
    finally:
        # 删除临时文件
        if os.path.exists(temp_output):
            os.remove(temp_output)
    return True


def add_watermark(input_video, watermark_image, output_video='output_with_watermark.mp4', position='bottom-right', offset_x=10, offset_y=10):
    """给视频添加水印，支持自定义位置和偏移量"""
    # 根据位置生成ffmpeg overlay参数
    position_map = {
        'top-left': '{0}:{1}'.format(offset_x, offset_y),
        'top-right': 'main_w-overlay_w-{0}:{1}'.format(offset_x, offset_y),
        'bottom-left': '{0}:main_h-overlay_h-{1}'.format(offset_x, offset_y),
        'bottom-right': 'main_w-overlay_w-{0}:main_h-overlay_h-{1}'.format(offset_x, offset_y)
    }

    overlay = position_map.get(position, position_map['bottom-right'])
    """给视频添加水印"""
    # 构建ffmpeg命令
    cmd = [
        'ffmpeg', '-i', input_video,
        '-i', watermark_image,
        '-filter_complex', 'overlay={0}'.format(overlay),  # 根据配置生成水印位置
        '-c:v', 'libx264',  # 复制视频流
        '-crf', '20',  # 设置质量因子，可以根据需要调整
        '-c:a', 'copy',  # 复制音频流
        '-y',
        output_video
    ]

    # 执行命令
    try:
        subprocess.check_output(cmd, stderr=subprocess.STDOUT, universal_newlines=True)
    except subprocess.CalledProcessError as e:
        print("Failed to add watermark: {}".format(e.output))
        return None
    print("Watermarked video saved to: {}".format(output_video))
    return output_video


def image_to_base64(image_path):
    """
    将图像文件转换为base64编码
    :param image_path: 图像文件路径
    :return: base64编码字符串
    """
    with open(image_path, "rb") as image_file:
        encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
    return encoded_string


def local_check_txt_nsfw(prompt):
    """使用OpenAI进行NSFW检测"""

    try:
        clientOpenAI = OpenAI()
        image_result = sr_utils.content_moderations(clientOpenAI, 3, "", prompt)
        print("image_result", image_result)
        return image_result
    except Exception as e:
        print("NSFW check failed: {}".format(str(e)))
        return False


def local_check_image_nsfw(image_path, prompt):
    """使用OpenAI进行NSFW检测"""

    try:
        # 将图像转换为base64
        image_base64 = image_to_base64(image_path)

        # 构造base64数据URL
        image_data_url = f"data:image/png;base64,{image_base64}"

        clientOpenAI = OpenAI()
        local_check_prompt_nsfw, code, err_message = sr_utils.content_moderations(clientOpenAI, 1, image_data_url, "")

        if code == 5:
            print(f"[ERR_MESSAGE] {err_message}")
            return True, 2, err_message
        return False, 2, f'7108,视频取帧GPT检测异常'
        # print("image_result", image_result)
        # return image_result
    except Exception as e:
        print("NSFW check failed: {}".format(str(e)))
        return False, 2, f'7108,视频取帧GPT检测异常' + str(e)


if __name__ == '__main__':
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='视频处理工具: 提取缩略图并添加水印')
    parser.add_argument('input_video', help='输入视频文件路径')
    parser.add_argument('--watermark-size', choices=['0', '360p', '480p', '540p', '720p', '768p', '1080p', '2k', '4k'],
                        help='水印尺寸，0表示不添加水印，其他值对应watermark文件夹中的PNG文件')
    parser.add_argument('--extract-thumbnail', action='store_true', help='是否提取缩略图，默认不提取')
    parser.add_argument('--nsfw-check', action='store_true', help='是否进行NSFW检测，默认不检测')
    parser.add_argument('--output-video', default='output_with_watermark.mp4', help='输出视频文件路径')
    parser.add_argument('--taskid', required=True, help='任务ID，用于创建帧存储目录')

    args = parser.parse_args()
    print(args)
    # 读取配置文件
    config = agent_config.SETTING_WATERMARK_THUMBNAIL_INFO

    # 处理水印路径
    watermark_path = None
    add_watermark_flag = False

    # 修改参数处理逻辑，当watermark-size为'0'时不添加水印
    if args.watermark_size and args.watermark_size != '0':
        # 修改为使用PNG格式的水印文件
        watermark_path = os.path.join('watermark', 'logo-{0}.png'.format(args.watermark_size))
        add_watermark_flag = True

        # 验证水印文件是否存在
        if add_watermark_flag and not os.path.exists(watermark_path):
            print('Error: Watermark file not found at {0}'.format(watermark_path))
            exit(11)

    # 设置默认参数
    thumbnail_time = config.get('thumbnail_time_seconds', 0)
    watermark_pos = config.get('watermark_position', 'bottom-right')
    offset_x = config.get('watermark_offset_x', 0)
    offset_y = config.get('watermark_offset_y', 0)

    # 新增: 提取指定时间点的帧并进行NSFW检测
    taskid = args.taskid  # 需要添加--taskid命令行参数
    task_dir = f"./download/{taskid}"
    frame_task_dir = os.path.join(task_dir, "watermark/image")
    os.makedirs(frame_task_dir, exist_ok=True)

    # 验证输入文件是否存在
    print("Input video file: {}".format(args.input_video))
    if not os.path.exists(args.input_video):
        print("Error: Input video file {} not found".format(args.input_video))
        exit(11)

    # 开始执行处理流程
    print("Starting video processing...")

    # 1、提取缩略图 - 只有当指定参数时才执行
    if args.extract_thumbnail:
        output_images = task_dir + "/_thumbnail.png"
        thumbnail = extract_thumbnail(args.input_video, output_image=output_images, time_second=thumbnail_time)
        if not thumbnail:
            exit(12)  # 提取缩略图失败

    # 2、添加水印 - 使用动态生成的watermark_path
    if add_watermark_flag:
        output_video = add_watermark(args.input_video, watermark_path, args.output_video, watermark_pos, offset_x, offset_y)
        if not output_video:
            exit(13)  # 添加水印失败

    # 3、提取帧并进行NSFW检测 - 只有当指定参数时才执行
    if args.nsfw_check:
        time_points = [1, 2.5, 3.5]
        for t in time_points:
            frame_path = frame_task_dir
            frame_path = os.path.join(frame_path, "frame_{}s.png".format(t))

            # 提取并压缩帧
            if not extract_frame(args.input_video, frame_path, t):
                print("Failed to extract frame at {} seconds".format(t))
                exit(14)  # 提取帧失败

            # NSFW检测
            # can also use base64 encoded image URLs
            # "url": "..."
            # False, 1, 'success'
            res, code, msg = local_check_image_nsfw(frame_path, "")
            if res:
                err_msg = f"{code},图片不合规: {frame_path}"
                print(err_msg)
                exit(15)  # NSFW检测失败


    # 示例 python add_watermark_service.py test_001_video.mp4 --output-video output.mp4 --taskid my_task_123 --watermark-size 720p --extract-thumbnail --nsfw-check
    # 示例 python add_watermark_service.py test_001_video.mp4 --output-video output.mp4 --taskid my_task_123